#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : knowledge.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/10 14:47
import json

import redis
from fastapi import Security, UploadFile, File, BackgroundTasks

from apps import logic
from apps.api.v1.knowledge_api import knowledge_router
from apps.api.v1.knowledge_api.schema import KnowledgeListRequests, KnowledgeListResponse, KnowledgeCreateRequests, \
    DefaultKnowledgeResponse, KnowledgeFileListRequests, KnowledgeFileListResponse, KnowledgeFileHandle, \
    CreateKnowledgeResponse, KnowledgeInfoResponse
from apps.api.v1.knowledge_api.util import handle_file, delete_knowledge_file, handle_file_
from apps.models import User
from apps.utils import get_current_user, calculate_md5, MinioUtil
from apps.utils.enumerate import KnowledgeFileStateEnum
from core.setting import settings


@knowledge_router.post('/list', name='知识库列表', response_model=KnowledgeListResponse)
async def knowledge_list(
        *,
        user: User = Security(get_current_user),
        query: KnowledgeListRequests
):
    count, data = await logic.knowledge_logic.knowledge_list(query=query, user=user)
    return {
        'code': 0,
        "message": "成功",
        "count": count,
        "data": data
    }


@knowledge_router.get('/all', name='获取全部的知识库', response_model=KnowledgeListResponse)
async def knowledge_all(
        *,
        user: User = Security(get_current_user),
):
    data = await logic.knowledge_logic.all_knowledge_object(user_id=user.id, is_delete=False)
    return {
        'code': 0,
        "message": "成功",
        "data": data
    }


@knowledge_router.post('/create', name='创建知识库', response_model=CreateKnowledgeResponse)
async def knowledge_create(
        *,
        user: User = Security(get_current_user),
        query: KnowledgeCreateRequests
):
    knowledge = await logic.knowledge_logic.knowledge_create(**query.dict(), user=user)
    if knowledge:
        return {
            'code': 0,
            "message": "成功",
            'data': knowledge.id
        }
    else:
        return {
            'code': -1,
            "message": "失败",
        }


@knowledge_router.post('/upload', name='知识库文件上传', response_model=DefaultKnowledgeResponse)
async def knowledge_file_upload(
        *,
        user: User = Security(get_current_user),
        knowledge_id: int,
        file: UploadFile = File(...)
):
    # 创建文件hash值
    file_hash = await calculate_md5(file=file)
    # 查询hash是否重复
    check = await logic.knowledge_file_logic.check_file_object(knowledge_id=knowledge_id, file_hash=file_hash)
    if check:
        return {
            'code': -1,
            'message': "文件重复"
        }
    # 获取文件类型
    file_type = file.filename.split('.')[-1]

    # 文件上传到minio
    try:
        minio_util = MinioUtil(
            secret_key=settings.MINIO_SECRET_KEY,
            access_key=settings.MINIO_ACCESS_KEY,
            endpoint=settings.MINIO_ENDPOINT,
        )
        bucket_name = settings.MINIO_BASE_BUCKET
        # 构造新的存储路劲
        object_name = f'knowledge/{user.id}/{knowledge_id}/{file_hash}.{file_type}'
        # 获取文件大小
        file.file.seek(0, 2)  # 将文件指针移动到文件末尾
        file_size = file.file.tell()  # 获取文件大小
        file.file.seek(0)  # 将文件指针重置到文件开头

        # 上传文件
        file_path = minio_util.put_object(bucket_name=bucket_name, object_name=object_name, data=file.file.read())
        await logic.knowledge_file_logic.crud.create(obj_in={
            'user_id': user.id,
            'file_name': file.filename,
            'file_size': file_size,
            'file_hash': file_hash,
            'file_address': file_path,
            'knowledge_id': knowledge_id,
            'state': KnowledgeFileStateEnum.waiting
        })
        return {
            'code': 0,
            "message": "成功"
        }
    except Exception as e:
        print(e)
        return {
            'code': -1,
            'message': "上传文件失败"
        }


@knowledge_router.post('/handle', name='处理知识库文件')
async def knowledge_file_handle(
        *,
        user: User = Security(get_current_user),
        query: KnowledgeFileHandle,
        background_tasks: BackgroundTasks
):
    if not await logic.knowledge_logic.knowledge_object(knowledge_id=query.knowledge_id, user=user):
        return {
            'code': -1,
            "message": "知识库不存在"
        }
    file_object = await logic.knowledge_file_logic.get_file_object(file_id=query.file_id,
                                                                   knowledge_id=query.knowledge_id)
    if not file_object:
        return {
            'code': -1,
            "message": "文件不存在"
        }
    # 查看文件状态
    if file_object.state == KnowledgeFileStateEnum.waiting:
        # 未处理的状态
        file_object.state = KnowledgeFileStateEnum.processing
        file_object.split_size = query.split_size  # 分片大小
        file_object.split_str = query.split_str  # 分片字符串
        if not query.file_extract_id:
            # 没有配置提取服务的话，使用默认的提取服务
            default_file_extract = await logic.file_extract_logic.get_default_file_extract()
            if not default_file_extract:
                return {
                    'code': -1,
                    'message': "请先配置提取服务"
                }
            query.file_extract_id = default_file_extract.id
        file_object.file_extract_id = query.file_extract_id  # 文件提取服务id

        # 修改处理方式为Celery
        try:
            pool = redis.ConnectionPool(host=settings.REDIS_HOST,
                                        port=settings.REDIS_PORT,
                                        db=settings.REDIS_DEFAULT_DB, decode_responses=True
                                        )
            r = redis.Redis(connection_pool=pool)
            # 数据存储
            _ = r.rpush(
                'chat_model:knowledge:file_handle', json.dumps(dict(
                    file_id=query.file_id,
                    knowledge_id=query.knowledge_id,
                    user_id=user.id
                ), ensure_ascii=False)
            )
        except Exception as e:
            print('redis存储异常', e)
            file_object.state = KnowledgeFileStateEnum.failed
        finally:
            r.close()
        await file_object.save()
        # await handle_file(
        #     file_id=query.file_id,
        #     knowledge_id=query.knowledge_id,
        #     user=user
        # )
        # background_tasks.add_task(
        #     handle_file,
        #     file_id=query.file_id,
        #     knowledge_id=query.knowledge_id,
        #     user=user
        # )
        return {
            'code': 0,
            "message": "成功"
        }

    elif file_object.state == KnowledgeFileStateEnum.processing:
        return {
            'code': -1,
            "message": "文件正在处理中"
        }
    elif file_object.state == KnowledgeFileStateEnum.success:
        # 处理成功
        return {
            'code': -1,
            "message": "重复处理成功文件"
        }
    else:
        # 处理失败
        return {
            'code': -1,
            "message": "文件处理失败"
        }


@knowledge_router.post('/file_list', name='获取文件列表', response_model=KnowledgeFileListResponse)
async def knowledge_file_list(
        *,
        user: User = Security(get_current_user),
        query: KnowledgeFileListRequests
):
    if not await logic.knowledge_logic.knowledge_object(knowledge_id=query.knowledge_id, user=user):
        return {
            'code': -1,
            "message": "知识库不存在"
        }
    count, data = await logic.knowledge_file_logic.knowledge_file_list(
        knowledge_id=query.knowledge_id, page=query.page, limit=query.limit, name=query.name
    )
    return {
        'code': 0,
        "message": "成功",
        "count": count,
        "data": data
    }


@knowledge_router.delete('/knowledge_file', name='删除知识库文件')
async def knowledge_file_delete(
        *,
        user: User = Security(get_current_user),
        knowledge_id: int,
        file_id: int
):
    knowledge = await logic.knowledge_logic.knowledge_object(knowledge_id=knowledge_id, user=user)
    if not knowledge:
        return {
            'code': -1,
            "message": "知识库不存在"
        }
    file_object = await logic.knowledge_file_logic.get_file_object(file_id=file_id, knowledge_id=knowledge_id)
    if not file_object:
        return {
            'code': -1,
            "message": "文件不存在"
        }

    if file_object.state not in [
        KnowledgeFileStateEnum.success, KnowledgeFileStateEnum.failed, KnowledgeFileStateEnum.waiting
    ]:
        return {
            'code': -1,
            'message': "文件当前状态不允许删除，请稍后再试"
        }
    # 删除文件
    minio_util = MinioUtil(
        secret_key=settings.MINIO_SECRET_KEY,
        access_key=settings.MINIO_ACCESS_KEY,
        endpoint=settings.MINIO_ENDPOINT,
    )
    bucket_name = settings.MINIO_BASE_BUCKET
    object_name = file_object.file_address
    minio_util.remove_object(bucket_name=bucket_name, object_name=object_name)
    try:
        if file_object.state == KnowledgeFileStateEnum.success:
            embedding_server = await logic.embedding_server_logic.embedding_server_info(
                embedding_server_id=knowledge.embedding_server_id)
            await delete_knowledge_file(table_name=embedding_server.table_name, knowledge_id=knowledge_id,
                                        file_id=file_id)
        file_object.is_delete = True
    except Exception as e:
        return {
            'code': -1,
            "message": f"删除失败：{e}"
        }
    await file_object.save()
    return {
        'code': 0,
        "message": "成功"
    }


@knowledge_router.delete('/delete', name='删除知识库')
async def knowledge_delete(
        *,
        user: User = Security(get_current_user),
        knowledge_id: int
):
    knowledge = await logic.knowledge_logic.knowledge_object(knowledge_id=knowledge_id, user=user)
    if not knowledge:
        return {
            'code': -1,
            "message": "知识库不存在"
        }
    # 获取全部的文件
    all_file = await logic.knowledge_file_logic.all_knowledge_file(knowledge_id=knowledge_id, user=user)
    if all_file:
        # 删除文件
        minio_util = MinioUtil(
            secret_key=settings.MINIO_SECRET_KEY,
            access_key=settings.MINIO_ACCESS_KEY,
            endpoint=settings.MINIO_ENDPOINT,
        )
        bucket_name = settings.MINIO_BASE_BUCKET
        for file in all_file:
            object_name = file.file_address
            minio_util.remove_object(bucket_name=bucket_name, object_name=object_name)
            file.is_delete = True
            await file.save()
        try:
            embedding_server = await logic.embedding_server_logic.embedding_server_info(
                embedding_server_id=knowledge.embedding_server_id)
            await delete_knowledge_file(table_name=embedding_server.table_name, knowledge_id=knowledge_id)
        except Exception as e:
            print('删除向量库内容失败,', e)
    knowledge.is_delete = True
    await knowledge.save()
    return {
        'code': 0,
        "message": "成功"
    }


@knowledge_router.get('/info', name='获取知识库详情', response_model=KnowledgeInfoResponse)
async def knowledge_info(
        *,
        user: User = Security(get_current_user),
        knowledge_id: int
):
    knowledge = await logic.knowledge_logic.knowledge_object(knowledge_id=knowledge_id, user=user)
    if knowledge:
        return {
            'code': 0,
            "message": "成功",
            "data": knowledge
        }
    else:
        return {
            'code': -1,
            "message": "知识库不存在"
        }
